'''
Created on Feb 25, 2012

@author: nzhao
'''
from const import *
from enthought.traits.api import *
from enthought.traits.ui.api import *
from enthought.enable.component_editor import ComponentEditor
from enthought.chaco.api import marker_trait, Plot, ArrayPlotData

class SingleNV(HasTraits):

#properties which describe the system
    nuc = Enum("None", "N14", "N15")
    magB = Range(0.0, 1000.0, 10.0)
    magTheta = Range(0.0,180.0, 0.0)
    magPhi = Range(0.0,360.0, 0.0)
    
    energy = Array(float,(1,3),[[1000.0*NV_ZFS, 1000.0*NV_ZFS, 0.0]])
    stateR = Array(float,(3,3),[[1.0, 0.0, 0.0],[0.0, 0.0, 1.0],[0.0, 1.0, 0.0]])
    stateI = Array(float,(3,3),[[0.0, 0.0, 0.0],[0.0, 0.0, 0.0],[0.0, 0.0, 0.0]])
    
    Svec = Array(float, (3,3), [[0., 0., 1.], [0., 0., -1.], [0., 0., 0.]])
    Ivec = Array(float, (3,3), [[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]])

    ch1enable  = Bool(True)
    ch1Phase = Float(0.0)
    ch1Freq = Float(2870.0)
    ch1Rabi = Float(0.0)
    
    ch2Phase = Float(0.0)
    ch2enable  = Bool(False)
    ch2Freq = Float(2870.0)
    ch2Rabi = Float(0.0)
    
    ch3Phase = Float(0.0)
    ch3enable  = Bool(False)
    ch3Freq = Float(2870.0)
    ch3Rabi = Float(0.0)
    
    ch4Phase = Float(0.0)
    ch4enable  = Bool(False)
    ch4Freq = Float(2870.0)
    ch4Rabi = Float(0.0)    
    
    rabi_xAxis = Str('np.linspace( 0, 20, 100)')
    rabi_calculate = Button()
    rabi_plot = Instance(Plot)
    
    odmr_xAxis= Str()
    
    cpmg_xAxis= Str()

#initialization for the class    
    def __init__(self):
        super(SingleNV, self).__init__()
        self.rabi_xlist = eval(self.rabi_xAxis)
        self.rabi_ylist = np.sin(self.rabi_xlist)
        
        plotdata = ArrayPlotData(x = self.rabi_xlist, y = self.rabi_ylist)
        default_plot = Plot(plotdata)
        self.renderer = default_plot.plot(("x", "y"), type="scatter", color="blue")[0]
        self.rabi_plot = default_plot

#GUI display structure    
    view = View(
                HSplit(
                    VGroup(
                        Group(
                              Item("nuc", label='nuclear spin'),
                              Item("magB", label='B [Gauss]'),
                              Item("magTheta", label='theta [degree]'),
                              Item("magPhi", label='phi [degree]'),
                              label="Single NV Parameters"
                        ),
                          
                        Group(
                              Item("energy", label='energy [MHz]', style='readonly'),
                              HGroup(
                                  Item("stateR", label='stateR',style='readonly'),
                                  Item("stateI", label='stateI',style='readonly'),
                                  show_border = True
                                  ),
                              label="Eigenenergy and Eigenstate"    
                        ),
                            
                        HGroup(
                              Item("Svec", label='S value',style='readonly'),
                              Item("Ivec", label='I value',style='readonly'),
                              show_border = True
                             ),
    
    
                        Group(                
                              HGroup(
                                     Item('ch1enable',  show_label=False),
                                     Item('ch1Phase', label="Phase", enabled_when='ch1enable'),
                                     Item('ch1Freq', label="Frequency", enabled_when='ch1enable'),
                                     Item('ch1Rabi', label="Rabi Frequency", enabled_when='ch1enable'),
                                     label='Channel 1'
                                     ),
                              HGroup(
                                   Item('ch2enable',  show_label=False),
                                   Item('ch2Phase', label="Phase", enabled_when='ch2enable'),
                                   Item('ch2Freq', label="Frequency", enabled_when='ch2enable'),
                                   Item('ch2Rabi', label="Rabi Frequency", enabled_when='ch2enable'),
                                   label='Channel 2'
                                   ),
                              HGroup(
                                   Item('ch3enable',  show_label=False),
                                   Item('ch3Phase', label="Phase", enabled_when='ch3enable'),
                                   Item('ch3Freq', label="Frequency", enabled_when='ch3enable'),
                                   Item('ch3Rabi', label="Rabi Frequency", enabled_when='ch3enable'),
                                   label='Channel 3'
                                   ),
                              HGroup(
                                   Item('ch4enable',  show_label=False),
                                   Item('ch4Phase', label="Phase", enabled_when='ch4enable'),
                                   Item('ch4Freq', label="Frequency", enabled_when='ch4enable'),
                                   Item('ch4Rabi', label="Rabi Frequency", enabled_when='ch4enable'),
                                   label='Channel 4'
                                   ),
                              label='MW'
                            ),scrollable = True
                    ),
                    Group(
                          Group(
                            HGroup(Item('rabi_xAxis', label='time [ns]', width=450),
                                    Item('rabi_calculate', show_label=False)
                                    ), 
                            Item('rabi_plot', editor=ComponentEditor(), show_label=False, width=650),
                            label="Rabi"
                          ),
                          Group(Item('odmr_xAxis', label='freq. [MHz]'), 
                                label="ODMR"),
                          Group(Item('cpmg_xAxis', label='time [ms]'), 
                                label="CPMG"),
                          layout='tabbed'
                          )
                ),
                title='Single NV',width = 1300, height = 600, 
                resizable = True
                )
    
#response functions for the user operations    
    def _rabi_calculate_fired(self):
        self.rabi_xlist = eval(self.rabi_xAxis)
        self.rabi_ylist = self.rabi_oscillation(self.rabi_xlist)
        
        plotdata = ArrayPlotData(x = self.rabi_xlist, y = self.rabi_ylist)
        plot = Plot(plotdata)
        self.renderer = plot.plot(("x", "y"), type="scatter", color="blue")[0]
        self.rabi_plot = plot
    
    def _nuc_changed(self):
        self.update_eigen_system()

    def _magB_changed(self):
        self.update_eigen_system()

    def _magTheta_changed(self):
        self.update_eigen_system()
    
    def _magPhi_changed(self):
        self.update_eigen_system()

#Kernel functions!!!! 
    def update_eigen_system(self):
        self.energy = [[NV_ZFS+self.magB,NV_ZFS-self.magB, 0.]]
        pass
    
    def rabi_oscillation(self, x):
        res = np.cos(x)
        return res

if __name__ == '__main__':
    singleNV=SingleNV()
    singleNV.configure_traits()